#include <functional>
#include <iostream>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

// 定义 Monoid 类型要求
template <typename T>
struct Monoid
{
    static T empty() { return T::empty(); }
    static T append(const T& a, const T& b) { return a + b; }
};

template <>
struct Monoid<std::string>
{
    static std::string empty() { return ""; }
    static std::string append(const std::string& a, const std::string& b) { return a + b; }
};

template <typename T>
struct Monoid<std::vector<T>>
{
    static std::vector<T> empty() { return {}; }
    static std::vector<T> append(const std::vector<T>& a, const std::vector<T>& b)
    {
        std::vector<T> result = a;
        result.insert(result.end(), b.begin(), b.end());
        return result;
    }
};

// 单元类型
struct Unit
{
    bool operator==(const Unit&) const { return true; }
};

// 复合 Monad 模板
template <typename R, typename W, typename S, typename A>
class RWSTMonad
{
public:
    using value_type = A;  // 用于类型推导

    // 计算类型: (环境, 状态) -> (值, 新状态, 日志)
    using Computation = std::function<std::tuple<A, S, W>(const R&, S)>;

    explicit RWSTMonad(Computation comp) : computation_(std::move(comp)) {}

    std::tuple<A, S, W> run(const R& env, S state) const { return computation_(env, state); }

    // 改进的 flatMap 实现
    template <typename F>
    auto flatMap(F f) const
    {
        // 使用 declval 推导返回类型
        using NextMonad = decltype(f(std::declval<A>()));
        using B = typename NextMonad::value_type;

        return RWSTMonad<R, W, S, B>(
            [this, f](const R& env, S state) -> std::tuple<B, S, W>
            {
                auto [a, newState, w1] = computation_(env, state);
                auto next = f(a);
                auto result = next.run(env, newState);
                auto [b, finalState, w2] = result;
                return {b, finalState, Monoid<W>::append(w1, w2)};
            });
    }

private:
    Computation computation_;
};

// ================== 基础操作 ==================

// 创建纯值
template <typename R, typename W, typename S, typename A>
RWSTMonad<R, W, S, A> pure(A value)
{
    return RWSTMonad<R, W, S, A>([value](const R&, S state) -> std::tuple<A, S, W> { return {value, state, Monoid<W>::empty()}; });
}

// 获取环境
template <typename R, typename W, typename S>
RWSTMonad<R, W, S, R> ask()
{
    return RWSTMonad<R, W, S, R>([](const R& env, S state) -> std::tuple<R, S, W> { return {env, state, Monoid<W>::empty()}; });
}

// 记录日志
template <typename R, typename W, typename S>
RWSTMonad<R, W, S, Unit> tell(const W& log)
{
    return RWSTMonad<R, W, S, Unit>([log](const R&, S state) -> std::tuple<Unit, S, W> { return {Unit{}, state, log}; });
}

// 获取状态
template <typename R, typename W, typename S>
RWSTMonad<R, W, S, S> get()
{
    return RWSTMonad<R, W, S, S>([](const R&, S state) -> std::tuple<S, S, W> { return {state, state, Monoid<W>::empty()}; });
}

// 设置状态
template <typename R, typename W, typename S>
RWSTMonad<R, W, S, Unit> put(S newState)
{
    return RWSTMonad<R, W, S, Unit>([newState](const R&, S) -> std::tuple<Unit, S, W> { return {Unit{}, newState, Monoid<W>::empty()}; });
}

// 修改状态
template <typename R, typename W, typename S, typename F>
auto modify(F transform)
{
    return get<R, W, S>().flatMap([transform](S state) { return put<R, W, S>(transform(state)); });
}

// ================== 辅助函数 ==================

// 打印结果
template <typename A, typename S, typename W>
void printResult(const std::string& title, const std::tuple<A, S, W>& result)
{
    auto [value, state, log] = result;
    std::cout << "=== " << title << " ===\n";
    std::cout << "Value: ";
    if constexpr (!std::is_same_v<A, Unit>)
    {
        std::cout << value;
    }
    else
    {
        std::cout << "<Unit>";
    }
    std::cout << "\nState: " << state;
    std::cout << "\nLog: \"" << log << "\"\n\n";
}
